import os
import random
import logging
from typing import List, Union
from config import Config, update_in_rule
import subprocess
from colorama import Fore, Back, Style

# 自定义日志格式化类，用于根据日志级别添加颜色
class ColoredFormatter(logging.Formatter):
    def __init__(self, fmt=None, datefmt=None):
        # 调用父类的初始化方法
        super().__init__(fmt, datefmt)

    def format(self, record):
        # 默认日志文字颜色为白色
        log_color = Fore.WHITE
        # 根据日志级别设置颜色
        if record.levelno == logging.DEBUG:  # 调试级别
            level_color = Fore.CYAN
        elif record.levelno == logging.INFO:  # 信息级别
            level_color = Fore.GREEN
        elif record.levelno == logging.WARNING:  # 警告级别
            level_color = Fore.YELLOW
        elif record.levelno == logging.ERROR:  # 错误级别
            level_color = Fore.RED
        elif record.levelno == logging.CRITICAL:  # 严重错误级别
            level_color = Fore.RED + Back.WHITE
        else:
            level_color = Fore.WHITE  # 其他级别默认为白色
        
        # 格式化日志级别部分，并添加颜色
        level_msg = f'[{record.levelname}]'
        colored_level_msg = level_color + level_msg + Style.RESET_ALL
        
        # 格式化日志消息部分，保持文字颜色为白色
        log_msg = super().format(record).replace(level_msg, colored_level_msg, 1)
        
        # 返回带有颜色级别和白色文字的日志消息
        return log_color + log_msg + Style.RESET_ALL

# 配置日志系统
formatter = ColoredFormatter(f'{Fore.WHITE}[%(levelname)s]{Style.RESET_ALL} {Fore.LIGHTBLUE_EX}%(message)s{Style.RESET_ALL}')
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logging.basicConfig(level=logging.DEBUG, handlers=[handler])

class TestGenerator:
    def __init__(self, data_dir: str = Config.PROBLEM_NAME):
        """
        初始化测试数据生成器
        :param data_dir: 测试数据存放目录
        """
        self.abs_main_dir = os.path.dirname(os.path.abspath(__file__))
        self.data_dir =os.path.join(self.abs_main_dir, data_dir)
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
    
    def compile_solution(self, solve_file: str, output_file: str) -> bool:
        """
        编译解答文件
        :return: 是否编译成功
        """
        try:
            if solve_file.endswith('.c'):
                subprocess.run(['gcc', "-Wno-return-type", solve_file, '-o', output_file], check=True)
            elif solve_file.endswith('.cpp'):
                subprocess.run(['g++', "-Wno-return-type", solve_file, '-o', output_file], check=True)
            return True
        except subprocess.CalledProcessError:
            logging.error("编译失败！")
            return False

    def generate_output(self, case_num: int, solve_file: str):
        """
        使用标准程序生成输出文件
        :param case_num: 测试样例编号
        """
        with open(os.path.join(self.abs_main_dir, f"{Config.PROBLEM_NAME}\\{case_num}.in"), 'r') as inf, \
                open(os.path.join(self.abs_main_dir, f"{Config.PROBLEM_NAME}\\{case_num}.out"), 'w') as outf:
            subprocess.run([f"{solve_file}"], 
                            stdin=inf, 
                            stdout=outf)
            
    def generate_number(self, value_rule: Union[List[int], int]) -> int:
        """
        根据值的类型生成随机数或者返回固定值
        :param value_rule: [最小值, 最大值] 或 固定值
        :return: 生成的随机数或固定值
        """
        if isinstance(value_rule, list):
            if len(value_rule) == 2:
                # 如果是范围规则，如 [min, max]
                return random.randint(value_rule[0], value_rule[1])
            elif len(value_rule) == 1:
                # 如果是固定值规则，如 [固定值]
                return value_rule[0]
            else:
                raise ValueError(f"规则格式错误：列表长度必须是 1 或 2，当前长度：{len(value_rule)}")
        else:
            raise ValueError(f"规则项必须是列表或整数，当前类型：{type(value_rule)}")

    def generate_line(self, rule: List[Union[List[int], int]]) -> str:
        """
        根据规则生成一行数据
        :param rule: 该行的生成规则，每个元素是一个[min, max]的范围，或一个固定值
        :return: 生成的数字串
        """
        numbers = [self.generate_number(range_rule) for range_rule in rule]
        return " ".join(map(str, numbers))
    
    def generate_single_test(self, case_num: int, rules: List[List[Union[List[int], int]]]) -> None:
        """
        生成单个测试用例
        :param case_num: 测试用例编号
        :param rules: 所有行的生成规则
        """
        in_path = os.path.join(self.data_dir, f"{case_num}.in")
        
        with open(in_path, 'w') as f:
            for rule in rules:
                line = self.generate_line(rule)
                f.write(line + '\n')
                
    def generate_test_cases(self, solve_file: str) -> None:
        """
        根据Config生成所有测试用例
        """
        # 得到绝对路径
        solve_file = os.path.join(self.abs_main_dir, solve_file)
        # 输出文件
        output_file = os.path.join(self.abs_main_dir, solve_file.split(".")[0])
        
        logging.info("开始编译答案文件...")
        if not self.compile_solution(solve_file=solve_file, output_file=output_file):
            return

        logging.info(f"开始生成 {Config.IN_NUM} 个测试用例...")

        for i in range(1, Config.IN_NUM + 1):
            if Config.IS_DIY:
                update_in_rule()
            logging.info(f"生成测试点 [{i}/{Config.IN_NUM}]")
            self.generate_single_test(i, Config.IN_RULE)
            self.generate_output(i, output_file)
        
        logging.info("测试用例生成完成！")
    
    def validate_config(self) -> bool:
        """
        验证配置是否合法
        :return: 配置是否合法
        """
        if not isinstance(Config.IN_NUM, int) or Config.IN_NUM <= 0:
            logging.error("错误：IN_NUM 必须是正整数")
            return False
            
        if not isinstance(Config.IN_RULE, list) or not Config.IN_RULE:
            logging.error("错误：IN_RULE 不能为空")
            return False
            
        for line_rule in Config.IN_RULE:
            if not isinstance(line_rule, list):
                logging.error("错误：每行规则必须是列表")
                return False
                
            for num_rule in line_rule:
                if isinstance(num_rule, list):
                    # 如果是列表，检查其长度
                    if len(num_rule) == 2:
                        if not isinstance(num_rule[0], int) or not isinstance(num_rule[1], int):
                            logging.error("错误：范围值必须是整数")
                            return False
                        if num_rule[0] > num_rule[1]:
                            logging.error("错误：范围的最小值不能大于最大值")
                            return False
                    elif len(num_rule) == 1:
                        # 如果是固定值，确保是整数
                        if not isinstance(num_rule[0], int):
                            logging.error("错误：固定值必须是整数")
                            return False
                    else:
                        logging.error(f"错误：规则格式不正确，列表长度必须为 1 或 2，当前列表长度：{len(num_rule)}")
                        return False
                elif isinstance(num_rule, int):
                    # 如果是整数，则认为是固定值
                    pass
                else:
                    logging.error(f"错误：规则项必须是列表或整数，当前类型：{type(num_rule)}")
                    return False
                    
        return True

def main():
    # 创建生成器实例
    generator = TestGenerator()

    diy = Config.IN_RULE
    
    # 验证配置
    if not generator.validate_config():
        return
        
    # 生成测试用例
    generator.generate_test_cases(solve_file=Config.SOLVE_FILE)
    
if __name__ == "__main__":
    main()